#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@CreateTime: 2025/10/11 19:56
@Author  : AnimateX
@Contact : animatex@163.com
@File    : timer_test.py
@License : Copyright © 2025 AnimateX. All rights reserved.
@version : timer_test_v2025_10_11

--------------------------------------------------------------

@Description: 
"""
import time
import functools
from typing import Callable, Any, Optional
from contextlib import contextmanager
import statistics
from collections import defaultdict


class Timer:
    """高性能计时器，支持多种使用方式和统计功能"""

    # 类级别的统计数据
    _stats = defaultdict(list)

    def __init__(self, name: str = "Timer", unit: str = "ms", precision: int = 4):
        """
        初始化计时器

        Args:
            name: 计时器名称
            unit: 时间单位 ('s', 'ms', 'us', 'ns')
            precision: 显示精度（小数点位数）
        """
        self.name = name
        self.unit = unit
        self.precision = precision
        self.start_time = None
        self.elapsed = None

        # 单位转换因子
        self.unit_factors = {
            's': 1,
            'ms': 1000,
            'us': 1000000,
            'ns': 1000000000
        }

        if unit not in self.unit_factors:
            raise ValueError(f"单位必须是: {list(self.unit_factors.keys())}")

    def start(self):
        """开始计时"""
        self.start_time = time.perf_counter()
        return self

    def stop(self):
        """停止计时并返回耗时"""
        if self.start_time is None:
            raise RuntimeError("计时器未启动，请先调用 start()")

        end_time = time.perf_counter()
        self.elapsed = (end_time - self.start_time) * self.unit_factors[self.unit]

        # 记录统计数据
        Timer._stats[self.name].append(self.elapsed)

        return self.elapsed

    def __enter__(self):
        """上下文管理器入口"""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.stop()
        print(f"[{self.name}] 耗时: {self.elapsed:.{self.precision}f} {self.unit}")
        return False

    @staticmethod
    def decorator(name: Optional[str] = None, unit: str = "ms", precision: int = 4,
                  print_result: bool = True, store_stats: bool = True):
        """
        装饰器模式计时

        Args:
            name: 函数名称（默认使用函数名）
            unit: 时间单位
            precision: 显示精度
            print_result: 是否打印结果
            store_stats: 是否存储统计数据
        """

        def decorator_timer(func: Callable) -> Callable:
            timer_name = name or func.__name__

            @functools.wraps(func)
            def wrapper(*args, **kwargs) -> Any:
                start = time.perf_counter()
                result = func(*args, **kwargs)
                end = time.perf_counter()

                elapsed = (end - start) * Timer(unit=unit).unit_factors[unit]

                if store_stats:
                    Timer._stats[timer_name].append(elapsed)

                if print_result:
                    print(f"[{timer_name}] 耗时: {elapsed:.{precision}f} {unit}")

                return result

            # 添加计时统计方法
            wrapper.get_stats = lambda: Timer.get_stats(timer_name)
            wrapper.reset_stats = lambda: Timer.reset_stats(timer_name)

            return wrapper

        return decorator_timer

    @staticmethod
    @contextmanager
    def measure(name: str = "Measurement", unit: str = "ms", precision: int = 4):
        """
        上下文管理器模式（不自动打印）

        Usage:
            with Timer.measure("task") as t:
                # your code
                pass
            print(f"耗时: {t.elapsed}")
        """
        timer = Timer(name=name, unit=unit, precision=precision)
        timer.start()
        try:
            yield timer
        finally:
            timer.stop()

    @staticmethod
    def get_stats(name: str) -> dict:
        """获取指定计时器的统计信息"""
        if name not in Timer._stats or not Timer._stats[name]:
            return {"error": "没有统计数据"}

        times = Timer._stats[name]
        return {
            "name": name,
            "count": len(times),
            "total": sum(times),
            "mean": statistics.mean(times),
            "median": statistics.median(times),
            "min": min(times),
            "max": max(times),
            "stdev": statistics.stdev(times) if len(times) > 1 else 0,
        }

    @staticmethod
    def get_all_stats() -> dict:
        """获取所有计时器的统计信息"""
        return {name: Timer.get_stats(name) for name in Timer._stats.keys()}

    @staticmethod
    def print_stats(name: Optional[str] = None, unit: str = "ms", precision: int = 4):
        """打印统计信息"""
        if name:
            stats = Timer.get_stats(name)
            if "error" in stats:
                print(stats["error"])
                return

            print(f"\n{'=' * 60}")
            print(f"统计信息: {name}")
            print(f"{'=' * 60}")
            print(f"调用次数: {stats['count']}")
            print(f"总耗时:   {stats['total']:.{precision}f} {unit}")
            print(f"平均耗时: {stats['mean']:.{precision}f} {unit}")
            print(f"中位数:   {stats['median']:.{precision}f} {unit}")
            print(f"最小值:   {stats['min']:.{precision}f} {unit}")
            print(f"最大值:   {stats['max']:.{precision}f} {unit}")
            print(f"标准差:   {stats['stdev']:.{precision}f} {unit}")
            print(f"{'=' * 60}\n")
        else:
            all_stats = Timer.get_all_stats()
            for timer_name in all_stats:
                Timer.print_stats(timer_name, unit, precision)

    @staticmethod
    def reset_stats(name: Optional[str] = None):
        """重置统计数据"""
        if name:
            Timer._stats[name] = []
        else:
            Timer._stats.clear()

    @staticmethod
    def compare_functions(*funcs, runs: int = 100, unit: str = "ms", precision: int = 4):
        """
        比较多个函数的性能

        Args:
            *funcs: 要比较的函数（可传入元组 (func, args, kwargs)）
            runs: 运行次数
            unit: 时间单位
            precision: 显示精度
        """
        results = []

        for item in funcs:
            if isinstance(item, tuple):
                func, args, kwargs = item[0], item[1] if len(item) > 1 else (), item[2] if len(item) > 2 else {}
            else:
                func, args, kwargs = item, (), {}

            times = []
            for _ in range(runs):
                start = time.perf_counter()
                func(*args, **kwargs)
                end = time.perf_counter()
                times.append((end - start) * Timer(unit=unit).unit_factors[unit])

            results.append({
                "name": func.__name__,
                "mean": statistics.mean(times),
                "median": statistics.median(times),
                "min": min(times),
                "max": max(times),
            })

        # 排序并打印结果
        results.sort(key=lambda x: x["mean"])

        print(f"\n{'=' * 70}")
        print(f"函数性能比较 ({runs} 次运行)")
        print(f"{'=' * 70}")
        print(f"{'排名':<6} {'函数名':<20} {'平均':<12} {'中位数':<12} {'最小':<12} {'最大':<12}")
        print(f"{'-' * 70}")

        for i, result in enumerate(results, 1):
            print(f"{i:<6} {result['name']:<20} "
                  f"{result['mean']:<12.{precision}f} "
                  f"{result['median']:<12.{precision}f} "
                  f"{result['min']:<12.{precision}f} "
                  f"{result['max']:<12.{precision}f} {unit}")
        print(f"{'=' * 70}\n")

        return results


# ============== 使用示例 ==============

if __name__ == "__main__":
    print("=" * 70)
    print("Timer 性能分析工具 - 使用示例")
    print("=" * 70)

    # 方式1: 装饰器模式
    print("\n1. 装饰器模式:")


    @Timer.decorator(unit="ms")
    def slow_function():
        time.sleep(0.1)
        return sum(range(10000))


    @Timer.decorator(name="快速函数", unit="us")
    def fast_function():
        return sum(range(100))


    slow_function()
    fast_function()

    # 方式2: 上下文管理器模式（自动打印）
    print("\n2. 上下文管理器模式（自动打印）:")
    with Timer("数据处理", unit="ms"):
        result = sum(range(1000000))

    # 方式3: 上下文管理器模式（手动获取时间）
    print("\n3. 上下文管理器模式（手动控制）:")
    with Timer.measure("自定义任务", unit="us") as t:
        data = [i ** 2 for i in range(10000)]
    print(f"自定义输出 - 任务耗时: {t.elapsed:.2f} {t.unit}")

    # 方式4: 手动启动/停止
    print("\n4. 手动启动/停止模式:")
    timer = Timer("手动计时", unit="ms")
    timer.start()
    time.sleep(0.05)
    elapsed = timer.stop()
    print(f"手动计时耗时: {elapsed:.4f} ms")

    # 方式5: 多次调用并查看统计
    print("\n5. 多次调用统计:")


    @Timer.decorator(name="重复任务", print_result=False)
    def repeated_task():
        time.sleep(0.01)
        return sum(range(1000))


    for _ in range(10):
        repeated_task()

    Timer.print_stats("重复任务")

    # 方式6: 函数性能比较
    print("\n6. 函数性能比较:")


    def method_a():
        return sum(range(10000))


    def method_b():
        return sum([i for i in range(10000)])


    def method_c():
        total = 0
        for i in range(10000):
            total += i
        return total


    Timer.compare_functions(method_a, method_b, method_c, runs=100, unit="us")

    # 方式7: 查看所有统计
    print("\n7. 查看所有累积的统计信息:")
    Timer.print_stats()  # 打印所有统计信息